'''
Module for communication across SALSA service platform

Based on AF_INET sockets

@author: myaut
'''

import threading
from threading import Thread, Event

import SocketServer
import socket
from Queue import Queue, Empty
from select import select

from salsa import trace_io, trace_io_obj, trace_exec_flow, do_trace, SALSAException, do_trace_exception

import json
import cPickle as pickle
import time

greeting_msg = 'Hello\n'

DEFERRED_TIMEOUT = 2
SOCKET_POLL_TIMEOUT = 0.05
QUEUE_WAIT_TIMEOUT = 0.05
MAX_FEEDER_SIZE = 32

def get_hostname(host):
    '''get_hostname(host) is helper function for master/slaves
    selects hostname if host==None (i.e. localhost) and makes extra
    checks if needed
    
    @todo deprecated, pass default hostnames to config'''
    if not host:
        return 'localhost'
    else:
        return host

class RawIOEncoderDecorer:
    def encode(self, obj):
        return str(obj)
    
    def decode(self, s):
        return str(s)

class JSONEncoderDecoder(json.JSONDecoder, json.JSONEncoder):
    def __init__(self):
        json.JSONDecoder.__init__(self)
        json.JSONEncoder.__init__(self)

class RPCEncoderDecoder:
    def encode(self, obj):
        return pickle.dumps(obj)
    
    def decode(self, s):
        return pickle.loads(s) 

class IOException(SALSAException):
    pass

class IOTransactionError(Exception):
    pass

class IOMasterDead:
    pass

class IORequestDead:
    def __init__(self, trace_buffer = None):
        self.trace_buffer = trace_buffer

class IOSlaveDead:
    def __init__(self, trace_buffer = None):
        self.trace_buffer = trace_buffer

class IOReply:
    tag = 'REP'
    
    def __init__(self, data):
        self.ready = True
        self.data = data
        
    def stop(self):
        pass
    
    def get(self):
        self.ready = False
        
        return self.data
    
    def __repr__(self):
        return self.__class__.__name__

class IOFinishReply(IOReply):
    tag = 'FIN'
    
    def __init__(self):
        pass
    
    def get(self):        
        return ''

class IODeferredReply(IOReply):
    '''Delays reply until data is ready.
    
    We need to block IORequestHandler thread instead of
    main service's loop, so we do it in get() routine'''
    tag = 'DFR'
    
    def __init__(self, data, event):
        self.data = data
        self.event = event
        
        self.is_done = False
        self.ready = True
    
    def wakeup(self):
        self.event.set()
    
    def stop(self):
        self.is_done = True
    
    def get(self):
        while not self.is_done:
            self.event.wait(DEFERRED_TIMEOUT)
            
            if self.event.is_set():
                break
        
        self.ready = False
        return self.data

class IOFeederReply(IOReply):
    '''Sends multiple replies until data is in generator
    Needed if data is too large (i.e. dump of repository DB,
    and we don't want read it into large buffer then send it.
    
    Feeder sends reply with small portions of data'''
    tag = 'FED'
    
    def __init__(self, iterable):
        self.generator = iter(iterable)
        self.ready = True
    
    def get(self):
        try:
            reply = next(self.generator)
        except StopIteration:
            self.ready = False
            raise
        
        return reply

class IOErrorReply(IOReply):
    '''Say, that transaction processing failed
    
    Same as IOReply, but have special tag
    that treated as Error in IOSlave.parse_reply'''
    tag = 'ERR'  

class IOBasicRequest:
    '''Abstract class for functions that shared between client-side and server-side objects of request
    
    Each request contains in_queue and on_queue and get()/put() pair of function:
    
                                 | <----+
                                 |      |
    get() <---  in_queue <---- recv()   |
                                 |      |
    put() ---> out_queue ----> send()   |        
                                 +------
    
    Contains low-level request object that represents socket or SocketServer's request handler 
    (actually may be any type of file like object that supports select(), recv() and send()
    '''
    trace = False
    
    def __init__(self):
        self.is_done = False
        
        self.in_queue = Queue()
        self.out_queue = Queue()
        
        self.ready = Event()
        
        self.request_fail = False
        
        self.trace_buffer = ''

    def __del__(self):
        if not self.is_done:
            self.stop()
  
    def set_request_object(self, request):
        '''Set low-level object that implements send/recv for request'''
        self.io_request = request
        
        # Cache fileno because socket.fileno() is pretty slow
        self.io_fileno = request.fileno()
    
    def do_trace(self, what, data):
        '''Tracing helper'''        
        tname = threading.current_thread().name
        
        trace_str = '%.12f @%12s [%5s] %s' % (time.time(), tname, what, repr(data))
        
        do_trace(trace_str)
        
        if self.trace:
            self.trace_buffer += trace_str + '\n';
            
    def awake(self):
        '''Awakes all threads that are waiting for data available
        using self.awake Event'''
        self.ready.set()
    
    # Interface for main Service thread
    def wait(self):
        '''Wait for available data'''
        if trace_io: self.do_trace('WAIT', None)
        
        while not self.is_done:
            self.ready.wait(QUEUE_WAIT_TIMEOUT)
            
            self.check_failures()
            
            if not self.in_queue.empty():
                break
        
        self.ready.clear()
    
    def get(self):
        '''Get data from in_queue. Uses check_failures() from derived class to check if something going wrong
        (i.e low-level request is dead). This version doesn't block thread calling it, only waits for QUEUE_WAIT_TIMEOUT
        
        @return data. If timed out, will return None'''
        try:                        
            if trace_io: self.do_trace('-GET', None)
            
            self.check_failures()
               
            data = self.in_queue.get(block = True, timeout = QUEUE_WAIT_TIMEOUT)
            
            self.check_failures()
            
            # if data:
            #    self.acknowledge()
            
            if trace_io: self.do_trace('GET', data)
            
            return data
        except Empty:
            pass
    
    def get_blocked(self):
        '''Like get(), gets data from in_queue, but blocks thread until data is not available
        
        @deprecated: Use wait() + get()'''
        while not self.is_done:
            data = self.get()
            
            if data != None:
                return data
    
    def put(self, reply):
        '''Puts response to out_queue
        
        @param reply - response data
        @param complete: - communication is finished. After sending last message, request will be closed'''
        if not self.is_done:
            self.out_queue.put(reply)
        
        if trace_io: self.do_trace('PUT', reply)
    
    # Internal request interface
    def poll(self, timeout = 0):
        '''Helper function that polls low-level request and checks if data available'''        
        rlist, wlist, xlist = select([self.io_fileno], [], [], timeout)
        
        return len(rlist) != 0 and rlist[0] == self.io_fileno
    
    def recv_data(self, buf_size):
        '''Receives data from socket using buf_size chunks.
        If no data available, ends, if socket is dead, sets is_done and request_fail flags
        and returns None
        
        @return: None if request is dead or received data'''
        data = ''
        
        if trace_io: self.do_trace('-RECV', 'Polling')
        
        # Wait for incoming data
        while not self.is_done:
            if self.poll(SOCKET_POLL_TIMEOUT):
                break
        else:
            return None
        
        if trace_io: self.do_trace('-RECV', 'Receving')
        
        # Receive data
        while True:
            chunk = self.io_request.recv(buf_size)
            
            if len(chunk) == 0:
                # Oops, socket is dead
                if trace_io: self.do_trace('FIN', self)
                
                self.is_done = True
                self.request_fail = True
                return None
                
            data += chunk
            
            if trace_io: self.do_trace('RECV', chunk)
            
            # No more data available, leave while
            if not self.poll():
                if trace_io: self.do_trace('+RECV', 'Done')
                break
            
        return data
    
    def send_data(self, data):
        if trace_io: self.do_trace('SEND', data)
            
        self.io_request.send(data)     
    
    def acknowledge(self):
        pass  

class IORequestHandler(SocketServer.BaseRequestHandler, IOBasicRequest):
    '''helper class that handles request produced by SocketServer'''
    trace = True
    
    def __init__(self, request, client_address, server):
        self.is_done = False
        
        self.was_dispatched = False
        
        self.reply = None
        self.finish_callback = None
        
        IOBasicRequest.__init__(self)  
        SocketServer.BaseRequestHandler.__init__(self, request, client_address, server)
        
        if trace_io_obj: 
            do_trace('IORqHandler %s created clnt %s -> %s:%s' % 
                    (repr(self), client_address, self.server.io_master.host, self.server.io_master.port) )
    
    def reg_finish_callback(self, callback):
        self.finish_callback = callback
    
    def send_reply(self, enc_dec):
        # If no reply exists (if we just created new object)
        # or it is not ready (all data in self.reply already processed)
        # extract new reply form out_queue and prepare it
        if not self.reply or not self.reply.ready:
            reply = self.out_queue.get()
        
            if not isinstance(reply, IOReply):
                reply = IOReply(reply)
            
            if isinstance(reply, IOFinishReply):
                self.is_done = True
            
            if trace_io: self.do_trace('REPLY', reply)
            
            self.reply = reply
        
        try:
            data = self.reply.get()
            send_data = self.reply.tag + enc_dec.encode(data)
            
            self.send_data(send_data)
        except StopIteration:        
            self.send_data('END')
    
    def handle(self):
        io_master = self.server.io_master
        
        self.set_request_object(self.request)
        self.send_data(greeting_msg)
        
        while not self.is_done:
            # Loop that sending/receiving data
            try:
                # read data from socket
                in_data = self.recv_data(io_master.buf_size)
                
                if not in_data:
                    break
                
                in_data = io_master.enc_dec.decode(in_data)
                
                # acknowledges are ignored
                if in_data != 'ACK':                       
                    # add self to dispatch queue in master
                    io_master.dispatch(self)
                    self.in_queue.put(in_data)
                    self.awake()
                
                # write data to socket                                          
                self.send_reply(io_master.enc_dec)
            except ValueError as e:
                do_trace_exception('IORequestHandler')
                
                self.is_done = True
                
                enc_dec_error = str(e)
                self.request.send(io_master.enc_dec.encode(enc_dec_error))
            except Exception as e:
                do_trace_exception('IORequestHandler')
                                
                self.request_fail = True
                self.is_done = True
    
    def finish(self):
        if self.finish_callback:
            self.finish_callback(self)
    
    def stop(self):
        if trace_io: do_trace('IORqHandler stopped %s' % (repr(self)) )
        
        if self.reply:
            self.reply.stop()
        
        self.is_done = True
    
    def check_failures(self):
        if self.request_fail:
            raise IORequestDead(self.trace_buffer)

class IOTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    daemon_threads = True

    allow_reuse_address = True
    
    def __init__(self, io_master):
        self.io_master = io_master
        
        SocketServer.TCPServer.__init__(self, (self.io_master.host, self.io_master.port), IORequestHandler)
    
class IOMaster:
    '''Master IO class (server in client-server terminology)
    
    Works it two modes:
        - DISPATCH_MESSAGE - request handler dispatches to disp_queue every time it 
          receives new message
        - DISPATCH_REQUEST - request handler dispatches only when first message arrives
    
    To set mode, write io_master.dispatch_mode = <> BEFORE starting it
    
    Wrapper for IOTCPServer'''
    DISPATCH_MESSAGE = 1
    DISPATCH_REQUEST = 2
    
    buf_size = 4096
    dispatch_mode = DISPATCH_MESSAGE
        
    def __init__(self, port, host=None, enc_dec=RawIOEncoderDecorer()):
        self.port = port
        self.host = get_hostname(host)
        
        self.enc_dec = enc_dec
        
        self.server = IOTCPServer(self)
        self.disp_queue = Queue()
        
        self.started = False
        self.is_done = False
    
    def dispatch(self, request):
        '''Dispatchs request. @see IOMaster'''
        if self.dispatch_mode == IOMaster.DISPATCH_MESSAGE \
                or self.dispatch_mode == IOMaster.DISPATCH_REQUEST and not request.was_dispatched:
            
            self.disp_queue.put(request)
            request.was_dispatched = True
               
    def start(self):
        if self.started:
            raise IOException('IOMaster.start() is not reenterable')
        
        self.started = True
        
        if trace_exec_flow: do_trace('IOMaster @%s:%d is started' % (self.host, self.port))
        
        self.server_thread = Thread(target = self.server.serve_forever, name = threading._newname('IOMasterThread-%d'))
        self.server_thread.setDaemon(True)
        
        self.server_thread.start()
    
    def stop(self):
        while not self.disp_queue.empty():
            rq = self.disp_queue.get()
            
            rq.stop()
        
        self.server.shutdown()
        
        self.is_done = True
        
        if trace_exec_flow: do_trace(trace_exec_flow, 'IOMaster @%s:%d is stopped' % (self.host, self.port))
    
    def get(self):
        while not self.is_done:
            try:                
                reply = self.disp_queue.get(block = True, timeout = QUEUE_WAIT_TIMEOUT)
                if trace_io: do_trace('IOMaster::GET Got request %s' % repr(reply))
                
                return reply
            except Empty:
                pass
        
        raise IOMasterDead()

class IOFeeder:
    '''Helper class that iterable, and provides a client-side interface for feeder
    
    Allows us to write this code:
    
    feeder = io_slave.get() # Master sends IOFeederReply
    
    for obj in feeder:
        ...1
    '''
    
    def __init__(self):
        self.stopped = False
        self.queue = Queue(maxsize = MAX_FEEDER_SIZE)
    
    def __repr__(self):
        return '<IOFeeder %d %d %s>' % (id(self), self.queue.qsize(), 
                                        'stopped' if self.stopped else '')
    
    def __iter__(self):
        return self
    
    def put(self, data):
        self.queue.put(data)
    
    def stop(self):
        if trace_exec_flow: do_trace('Feeder %s is stopped' % self)
        
        self.stopped = True
    
    def next(self):       
        while True:
            try:
                return self.queue.get(block = True, timeout = QUEUE_WAIT_TIMEOUT)
            except Empty:
                if self.stopped:
                    raise StopIteration()

class IOSlave(IOBasicRequest):
    get_timeout = 1
    buf_size = 4096
    trace = True
    
    def __init__(self, port, host=None, enc_dec=RawIOEncoderDecorer()):
        self.port = port 
        self.host = get_hostname(host)
        
        self.enc_dec = enc_dec
        
        self.started = False
        self.start_event = Event()
        
        self.feeder = None
        
        self.socket_failure = None
        self.tx_failure = None
        
        if trace_io_obj: do_trace('IOSlave created %s -> %s:%d' % (repr(self), self.host, self.port) )
        
        IOBasicRequest.__init__(self)
    
    def start(self):
        self.init_socket()
        
        self.socket_thread = Thread(target = self._start_socket, name = threading._newname('IOSlaveThread-%d'))
        self.socket_thread.start()
        
        if trace_exec_flow: do_trace('IOSlave ->%s:%d is started' % (self.host, self.port))
    
    def stop(self):
        # Stop thread (if it is not blocked on recv)
        self.is_done = True
        
        if trace_io_obj: do_trace('IOSlave stopped %s -> %s:%d' % (repr(self), self.host, self.port) )
        if trace_exec_flow: do_trace('IOSlave ->%s:%d is stopped' % (self.host, self.port))
    
    def init_socket(self):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        
        try:
            self.socket.connect((self.host, self.port))
            
        except socket.error, (errno, message):
            raise IOException('Socket error %d: %s' % (errno, message))
    
    def put(self, data):
        self.start_event.wait(3)
        
        if not self.started:
            raise IOException('slave is not started (may be dead)')
        
        return IOBasicRequest.put(self, data)
    
    def parse_reply(self, reply):
        if not reply or len(reply) < 3:
            return None
        
        tag = reply[:3]
        data = None
        
        if len(reply) > 3:
            data = self.enc_dec.decode(reply[3:])
        
        if trace_io: self.do_trace('TAG', tag)
        
        if tag == IOErrorReply.tag:
            self.tx_failure = data
            self.awake()
        elif tag == IOFeederReply.tag:
            if not self.feeder:
                self.feeder = IOFeeder()
                self.in_queue.put(self.feeder)
            
            if trace_io: self.do_trace('FED', data)
            
            self.feeder.put(data)
            self.acknowledge()
        elif tag == 'END':
            if self.feeder:
                self.feeder.stop()
            else:
                # If no data available, IOFeederReply will immediately give us END
                # without any "FED" tags
                # Add to in_queue empty list, so upper layer will think that it is empty IOFeeder
                
                self.in_queue.put([]) 
                
            self.feeder = None
        elif tag == 'FIN':
            self.is_done = True
            return None
        else:
            return data
        
    def acknowledge(self):
        if self.feeder:
            self.out_queue.put('ACK')
    
    def _start_socket(self):
        if self.started:
            self._fail('IOSlave._start_socket() is not reenterable')
        
        self.set_request_object(self.socket)
        
        self.started = True
        self.start_event.set()
        
        self.recv_greeting()
        
        while not self.is_done:           
            try:
                out_data = self.out_queue.get(block = True, timeout = self.get_timeout)
            except Empty:
                continue
                
            self.send_data(self.enc_dec.encode(out_data))          
            
            reply = self.recv_data(self.buf_size)
            
            in_data = self.parse_reply(reply)
            
            if in_data != None:
                self.in_queue.put(in_data)
                self.awake()
            
        self.socket.close()
    
    def _fail(self, message):
        self.socket_failure = message
        raise IOException(message)
    
    def recv_greeting(self):
        # Wait 5s for greeting
        if not self.poll(5):
            self._fail('Master not responding')
        
        try:
            data = self.socket.recv(self.buf_size)
            
            if trace_io: self.do_trace('GREET', data)
            
            if data != greeting_msg:
                self._fail('Wrong greeting')
        except socket.error, (errno, message):
            self._fail('Socket error %d: %s' % (errno, message))
     
    def check_failures(self):
        if self.socket_failure:
            raise IOException(self.socket_failure)
        
        if self.tx_failure:
            tx_failure = self.tx_failure
            self.tx_failure = None
            
            raise IOTransactionError(tx_failure)
        
        if self.request_fail or self.is_done:
            raise IOSlaveDead(self.trace_buffer)

class IOSlaveAdapter:
    def __init__(self, io_slave):
        self.io_slave = io_slave

class IOSlaveCommunicator(IOSlaveAdapter):
    def communicate(self, arg):
        self.io_slave.put(arg)
        
        self.io_slave.wait()
        
        return self.io_slave.get()